#!/usr/bin/env python

from collections import defaultdict
import cPickle
import glob
import os
import optparse
import sys

import tables

import charmicat


class EfficiencyMonitor(charmicat.MPMonitor):
    def __init__(self, nproc, nfiles, logdir):
        super(EfficiencyMonitor, self).__init__(nproc, logdir=logdir, interval=2)
        self.files_total = nfiles

        self.done_glob_pattern = '*.pkl'
        
        self.final_effd = defaultdict(lambda: 0)

    def process_logs(self):
        self.files_processed = 0

        for path in self.logfiles:
            fp = open(path)
            try:
                nfile = int(fp.readlines()[-1])
            except ValueError:
                nfile = 0

            self.files_processed += int(nfile)
            
        sys.stdout.write("{0}/{1} files processed (polls={2})\r".format(self.files_processed, 
                                                                        self.files_total,
                                                                        self.npolls))
        sys.stdout.flush()
        
    def finalize(self):
        print "{0}/{1} files processed (polls={2})\r".format(self.files_total,
                                                             self.files_total,
                                                             self.npolls)

        effd_paths = glob.glob(os.path.join(self.logdir, self.done_glob_pattern))

        for p in effd_paths:
            with open(p) as fp:
                effd = cPickle.load(fp)
                for k in effd:
                    self.final_effd[k] += effd[k]
        
        


class EfficiencyWorker(charmicat.MPWorker):
    def __init__(self, logdir, args):
        super(EfficiencyWorker, self).__init__(logdir, args)
        
#        self.logfile_tag = args[0].name
#        self.errfile_tag = args[0].name
#        self.donefile_tag = args[0].name

    def work(self, cut_list, input_files, opts):     
 #       self.logfile_tag = cut.name

        self.npass = defaultdict(lambda: 0)
        self.ntot  = 0
        self.nfile = len(input_files)

        self.ntot = 0
        all_condvars = {}
        for ip, path in enumerate(input_files):
            with tables.openFile(path) as fp:
                self.ntot  += len(fp.root.nt)
                for cut in cut_list:
                    all_condvars.update(cut.condvars)
                    try:
                        self.npass[cut.name] += len(fp.root.nt.getWhereList(cut.pos_str, condvars=cut.condvars))
                    except NameError, e:
                        raise NameError(str(e) + ' ' + " on cut `{0}'".format(cut.name))
                
                self.npass['__all__'] += len(fp.root.nt.getWhereList(charmicat.cut_union_str(cut_list), 
                                                                     condvars=all_condvars))
        
            self.log(ip, self.nfile)
        
        self.log(self.nfile, self.nfile)

        self.npass['__total__'] = self.ntot
        
    def completed(self):
        path = os.path.join(self.logdir, '{0}.pkl'.format(os.getpid()))
        with open(path, 'w') as fp:
            cPickle.dump(dict(self.npass), fp)


class BaseEffColumn(object):
    def __init__(self, name, title, width, fmttype):
        self.name    = name
        self.title   = title
        self.width   = width
        self.fmttype = fmttype

        self.header = "{0:<{1}}".format(self.title, self.width)

    def get_field(self, *args):
        return "{0:<{1}{2}}".format(self.value(*args),
                                    self.width, self.fmttype)

    def value(self, *args):
        return 0


class IndexEffColumn(BaseEffColumn):
    def value(self, *args):
        return args[1].index(args[0])
    
        
class AttrEffColumn(BaseEffColumn):
    def __init__(self, name, title, width, fmttype, attr):
        super(AttrEffColumn, self).__init__(name, title, width, fmttype)

        self.attr = attr
        
    def value(self, *args):
        return getattr(args[0], self.attr)


class DictEffColumn(BaseEffColumn):
    def __init__(self, name, title, width, fmttype, key):
        super(DictEffColumn, self).__init__(name, title, width, fmttype)

        self.key = key
        
    def value(self, *args):
        return args[2][self.key]

class Arg0AttrDictEffColumn(BaseEffColumn):
    def __init__(self, name, title, width, fmttype, attr):
        super(Arg0AttrDictEffColumn, self).__init__(name, title, width, fmttype)

        self.attr = attr
        
    def value(self, *args):
        return args[2][getattr(args[0], self.attr)]


class PctEffColumn(BaseEffColumn):        
    def value(self, *args):
        return 100*(float(args[2][args[0].name]) / args[2]['__total__'])



_g_allcols = dict([(col.name, col) for col in [IndexEffColumn('index', 'Cut #', '3', 'd'),
                                               AttrEffColumn('name', 'Cut name', '25', 's', 'name'),
                                               AttrEffColumn('latex', 'LaTeX', '50', 's', 'latex'),
                                               AttrEffColumn('desc', 'Description', '70', 's', 'desc'),
                                               AttrEffColumn('pos_str', 'Cut in code', '70', 's', 'pos_str'),
                                               DictEffColumn('in', '# in', '12', 'd', '__total__'),
                                               Arg0AttrDictEffColumn('out', '# out', '12', 'd', 'name'),
                                               PctEffColumn('pct', "% pass", '6', ".3f"),
                                               ]])
                                                                                              

def colfcn_index(cut, cutlist, effd):
    return cutlist.index(cut)

def colfcn_name(cut, cutlist, effd):
    return cut.name

def colfcn_latex(cut, cutlist, effd):
    return cut.latex

def colfcn_desc(cut, cutlist, effd):
    return cut.desc

def colfcn_pos_str(cut, cutlist, effd):
    return cut.pos_str

def colfcn_in(cut, cutlist, effd):
    return effd['__total__']

def colfcn_out(cut, cutlist, effd):
    return effd[cut.name]

def colfcn_pct(cut, cutlist, effd):
    return 100*(float(effd[cut.name]) / effd['__total__'])


allowed_cols = tuple("index name latex desc pos_str in out pct".split())
col_titles   = dict(zip(allowed_cols, "Cut #,Cut name,LaTeX,Description,Cut in code,# in,# passed,% pass".split(',')))
col_widths   = dict(zip(allowed_cols, "3 25 50 70 70 12 12 6".split()))
col_headers  = dict(zip(allowed_cols, ["{0:<{1}}".format(col_titles[col], col_widths[col]) for col in allowed_cols]))
col_fmttypes = dict(zip(allowed_cols, "d s s s s d d .3f".split()))
col_fcns     = dict(zip(allowed_cols, [colfcn_index, 
                                       colfcn_name,
                                       colfcn_latex,
                                       colfcn_desc,
                                       colfcn_pos_str,
                                       colfcn_in,
                                       colfcn_out,
                                       colfcn_pct]))

#        for c in cut_list:
#            fpout.write(rowstr.format(c.name, c.pos_str, 
#                                      monitor.final_effd['__total__'], 
#                                      monitor.final_effd[c.name],
#                                      100*(float(monitor.final_effd[c.name]) / monitor.final_effd['__total__']),
#                                      max_name_len=max_name_len,
#                                      max_code_len=max_code_len))


def main():
    """
    ccat_efficiencies.py [opts] file1.h5 [file2.h5 [...]]

    Calculate the efficiencies for the default cutset for the provided files.
    """

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option("-c", "--cutset", dest="cutsets", action='append', default=[],
                  help="Apply cutset NAME (can be used more than once)")
    op.add_option("-m", "--module", dest='module_name', default=None, 
                  help="Use module MODULE instead of inferring from input", metavar='MODULE')
    op.add_option("-l", "--logdir", dest='logdir', default=None,
                  help='Specify directory PATH for MP logging', metavar='PATH')
    op.add_option("-o", "--output", dest="output", default=None,
                  help="Place output in directory NAME", metavar='NAME')
    op.add_option("-r", "--replace", action="append", dest="replace", default=[], nargs=2,
                  help="Replace in cutset cut named CUT1 with cut named CUT2")
    op.add_option("-s", dest="list", action='append', default=[],
                  help="Apply the cut NAME (can be used more than once)")
    op.add_option("-t", "--title", dest='title', default=None, 
                  help='Set title in output file to TITLE', metavar='TITLE')
    op.add_option("-u", "--union", action="store_true", dest="union", default=False,
                  help="Write output for the union of all cuts specified")
    op.add_option("-x", dest='exclude', action='append', default=[],
                  help="Exclude cut NAME (can be repeated)", metavar='NAME')
    op.add_option("--cols", dest="cols", default='name,in,out,pct',
                  help="Specify output columns in format COL1,COL2,COL3", metavar="COL1,COL2,COL3")
    op.add_option("--cpus", dest="ncpu", type="int", default=charmicat.DEFAULT_NCPU,
                  help="Spawn no more than N threads", metavar="N")
    op.add_option("--force-log", dest='force_log', default=False, action='store_true',
                  help='Force reuse of currently-existing log directory')
    op.add_option("--start-on", dest="start_on", default=None,
                  help="Apply only cuts downstream (inclusive) of NAME", metavar='NAME')
    
    opts, args = op.parse_args()

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise charmicat.BadCommandLineError(msg, -1)

    if not opts.module_name:
        opts.module_name = charmicat.module_name_from_paths(args)
    
    try:
        module = __import__(opts.module_name)           
    except ImportError:
        msg = "Error importing module `{0}'".format(opts.module_name)
        raise charmicat.BadCommandLineError(msg, -2)

    cut_list = charmicat.build_cut_list(module, opts)

    args = [os.path.abspath(arg) for arg in args]
    
    if not opts.output:
        opts.output = '+'.join(opts.cutsets) + '.efficiencies.txt'

    if opts.start_on:
        start_on_reached = False
    
    if len(args) < opts.ncpu:
        opts.ncpu = len(args)

    opts.cols = opts.cols.split(',')

    for col in opts.cols:
        if col not in _g_allcols:
            msg = "Unknown column `{0}' requested; known cols are {1}".format(col, sorted(_g_allcols.keys()))
            raise BadCommandLineError(msg, -3)

    opts.cols = [_g_allcols[cn] for cn in opts.cols]

    input_lists = charmicat.n_groups(opts.ncpu, args)
        
    monitor = EfficiencyMonitor(len(input_lists), len(args), opts.logdir)                
    workers = [EfficiencyWorker(monitor.logdir, (cut_list, il, opts)) for il in input_lists]

    charmicat.mp_deploy_and_monitor(monitor, workers, ncpu=opts.ncpu)

#    max_field_widths = {}
    
#    for f in 'name pos_str desc'.split():
#        max_field_widths[f] = max(max([len(getattr(cut, f)) for cut in cut_list]),
#                                  len(col_titles[f]))
        
    if opts.title is None:
        if hasattr(module, 'channel_latex'):
            opts.title = module.channel_latex + " efficiencies for `{0}' cuts:".format('+'.join(opts.cutsets))
        else:
            opts.title = opts.module_name + " efficiencies for `{0}' cuts:".format('+'.join(opts.cutsets))

    lines = [opts.title, '\n']
    lines.append(' '.join([col.header for col in opts.cols]))

    for cut in cut_list:
        row = []
        for col in opts.cols:
            row.append(col.get_field(cut, cut_list, monitor.final_effd))
        
        lines.append(' '.join(row))


    with open(opts.output, 'w') as fpout:
        fpout.write('\n'.join(lines))
        fpout.write('\n')

        fpout.write('\n')
        totstr = "Total efficiency: {0} of {1} candidates pass ({2:.3f} %)\n"
        fpout.write(totstr.format(monitor.final_effd['__all__'],
                                  monitor.final_effd['__total__'],
                                  100*(float(monitor.final_effd['__all__']) / monitor.final_effd['__total__']),
                                  ))
        

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)

